# -*- coding: utf-8 -*-
# @日期    : 2021/11/28 15:03
# @作者  : 万方名
# @FileName: main.py


import xgboost as xgb

xgb_train = xgb.DMatrix('../data/MQ2008/Fold1/mq2008.train')
xgb_test = xgb.DMatrix('../data/MQ2008/Fold1/mq2008.test')
xgb_vali = xgb.DMatrix('../data/MQ2008/Fold1/mq2008.vali')

# train model
params = {
    'objective': 'rank:pairwise',  # 学习目标：排序问题
    'gamma': 1.0,
    'min_child_weight': 0.1,  # 叶子结点最小样本权重和，若节点分裂导致叶子结点的样本权重和小于该值则节点不进行分裂
    'eta': 0.1,  # 学习率
    'max_depth': 6  # 决策树分裂的最大深度
}

num_round = 50
watchlist = [(xgb_train, 'train'), (xgb_vali, 'vail')]

model = xgb.train(params, xgb_train, num_round, watchlist)

pred = model.predict(xgb_test)
print(pred)